#!/usr/bin/python

from multiprocessing import Lock
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from threading import Thread

import random
import sys
import time

from util.time import get_interval_stamp, print_interval


# Reference: https://blog.csdn.net/u013210620/article/details/78687028
class FreqControl(Thread):
    def __init__(self, thread_count, min_sleep_time, max_sleep_time):
        super().__init__()
        self.pool = ThreadPoolExecutor(thread_count)
        self.min_time = min_sleep_time
        self.max_time = max_sleep_time
        self.need_to_sleep = (self.min_time + self.max_time > 99)
        self.total_job = 0
        self.finish_job = 0
        self.finish_dict = {}
        self.base = time.time()
        self.finish_count_lock = Lock()
        self.console_lock = Lock()

    def __sync_accumulate_count(self, job_id):
        self.finish_count_lock.acquire()
        try:
            self.finish_job += 1
            self.finish_dict[job_id] = self.finish_job
        finally:
            self.finish_count_lock.release()

    def __sync_get_count(self, job_id):
        self.finish_count_lock.acquire()
        try:
            val = self.finish_dict[job_id]
        finally:
            self.finish_count_lock.release()
        if val is None:
            raise Exception("无法获取完成任务数量")
        else:
            return val

    def __sync_print(self, text, auto_new_line=False):
        self.console_lock.acquire()
        try:
            if auto_new_line:
                print("\r{}".format(text), end="\n")
            else:
                print("\r{}".format(text), end="")
            sys.stdout.flush()
        finally:
            self.console_lock.release()

    def __proxy_method(self, **kwargs):
        method = kwargs['method']
        job_id = kwargs['id']
        arg = kwargs['args']
        ret = method(args=arg)
        ret['id'] = job_id
        if self.need_to_sleep:
            time.sleep(random.uniform(self.min_time, self.max_time) / 1000.0)
        self.__sync_accumulate_count(job_id)
        return ret

    def wait_finish(self):
        self.pool.shutdown(wait=True)

    def show_progress(self, job_id, text, auto_new_line=False):
        finished_count = self.__sync_get_count(job_id)
        percentage = float(finished_count) / self.total_job
        velocity = finished_count / get_interval_stamp(self.base)
        remaining_time = print_interval(int((self.total_job - finished_count) / (velocity if velocity > 0 else 1)))
        length = 30
        dash_count = int(length * percentage)
        arrow_count = 1 if dash_count > 0 else 0
        dash_count -= arrow_count
        progress = "\r%.2f%% %.3f 任务/秒 剩余时间：%s[%s%s%s] 完成进度 %d / %d 查询任务! " % (100.0 * percentage, velocity, remaining_time, "-" * dash_count, ">" * arrow_count, " " * (length - dash_count - arrow_count), finished_count, self.total_job)
        text = '' if text is None else text
        self.__sync_print(text if auto_new_line else progress + text, auto_new_line)

    def run_jobs(self, method, job_ids, args=None, callback=None, **kwargs):
        self.base = time.time()
        if args is not None:
            # 通过args传入list类型参数
            self.total_job = len(args)
            if callback is None:
                for job_id, arg in zip(job_ids, args):
                    self.pool.submit(self.__proxy_method, method=method, args=arg, id=job_id)
            else:
                for job_id, arg in zip(job_ids, args):
                    self.pool.submit(self.__proxy_method, method=method, args=arg, id=job_id).add_done_callback(callback)
        elif len(kwargs) > 0:
            # 通过直接定义参数名和数值传入参数
            param_count = len(kwargs)
            param_name = []
            param_values = []
            job_count = len(job_ids)
            # 检查各个参数对应的任务数量，预处理
            for name, values in kwargs.items():
                temp_count = len(values)
                if job_count != temp_count:
                    raise Exception("不同参数对应的任务数不一致")
                param_name.append(name)
                param_values.append(list(values))
            self.total_job = job_count
            # 生成每个任务对应的参数
            params = []
            job_id_list = list(job_ids)
            for i in range(job_count):
                param = {'id': job_id_list[i]}
                for j in range(param_count):
                    param[param_name[j]] = param_values[j][i]
                params.append(param)
            if callback is None:
                for param in params:
                    self.pool.submit(self.__proxy_method, method=method, args=param, id=param['id'])
            else:
                for param in params:
                    self.pool.submit(self.__proxy_method, method=method, args=param, id=param['id']).add_done_callback(callback)

    @staticmethod
    def get_result(obj):
        return obj.result()


if __name__ == '__main__':
    def show(*args, **kwargs):
        if len(kwargs) > 0:
            num = kwargs['args']['num']
            base = kwargs['args']['base']
        elif len(args) > 0:
            num = args[0]
            base = args[1]
        else:
            raise Exception("参数为空")
        print("Running job " + str(num) + " at: " + str(get_interval_stamp(base)))
        return {'num': num, 'base': base}

    def callback_func(obj):
        res = obj.result()
        print("Finish  job " + str(res['num']) + " at: " + str(get_interval_stamp(res['base'])))

    lst = range(30)
    fc = FreqControl(2, 300, 700)
    base = time.time()
    bases = [base for num in lst]
    # 方式1：通过直接定义参数名和数值传入参数
    # fc.run_jobs(show, num=lst, base=bases, callback=callback_func)
    # 方式2：通过args传入list类型参数
    fc.run_jobs(show, args=[{"num": num, "base": base} for num in lst], callback=callback_func)
    fc.wait_finish()
    print("All done at: " + str(get_interval_stamp(base)))
